{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c2b24b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "from scipy.io import loadmat\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.model_selection import train_test_split\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "# -------------------------- 数据加载与预处理（严格遵循网页逻辑） --------------------------\n",
    "# 1. 定义故障标签（与网页完全一致，10种故障类型）\n",
    "FAULT_LABEL_DICT = {\n",
    "    '97': 0,    # 正常状态（对应X097.mat）\n",
    "    '105': 1,   # 内圈故障\n",
    "    '118': 2,   # 外圈故障（6点钟位置）\n",
    "    '130': 3,   # 滚动体故障\n",
    "    '169': 4,   # 内圈故障（更高转速）\n",
    "    '185': 5,   # 外圈故障（6点钟，更高转速）\n",
    "    '197': 6,   # 滚动体故障（更高转速）\n",
    "    '209': 7,   # 内圈故障（另一种载荷）\n",
    "    '222': 8,   # 外圈故障（12点钟）\n",
    "    '234': 9    # 滚动体故障（另一种载荷）\n",
    "}\n",
    "AXIS = '_DE_time'  # 仅使用驱动端数据\n",
    "\n",
    "def load_cwru_data(data_path, time_steps=1024, window=128):\n",
    "    \"\"\"\n",
    "    按网页逻辑加载CWRU数据集，支持10类故障，滑动窗口切片\n",
    "    time_steps: 样本长度（网页中为1024）\n",
    "    window: 相邻样本重叠点数（网页中为128，步长=time_steps - window）\n",
    "    \"\"\"\n",
    "    X, y = [], []\n",
    "    for fault_type, label in FAULT_LABEL_DICT.items():\n",
    "        # 构建文件名和数据键名（处理正常状态的特殊键名X097_DE_time）\n",
    "        mat_file = os.path.join(data_path, f'X{fault_type}.mat')\n",
    "        if fault_type == '97':\n",
    "            data_key = f'X0{fault_type}{AXIS}'\n",
    "        else:\n",
    "            data_key = f'X{fault_type}{AXIS}'\n",
    "        \n",
    "        # 加载mat文件并提取数据\n",
    "        if os.path.exists(mat_file):\n",
    "            mat_data = loadmat(mat_file)[data_key].flatten()\n",
    "            # 滑动窗口切片（步长=time_steps - window，与网页一致）\n",
    "            for i in range(0, len(mat_data) - time_steps + 1, time_steps - window):\n",
    "                X.append(mat_data[i:i+time_steps])\n",
    "                y.append(label)\n",
    "    \n",
    "    return np.array(X, dtype=np.float32), np.array(y, dtype=np.int64)\n",
    "\n",
    "# 2. 自定义数据集类（匹配网页数据格式）\n",
    "class CWRUDataset(Dataset):\n",
    "    def __init__(self, X, y, mode='train', val_rate=0.3, test_rate=0.5, is_train=True):\n",
    "        \"\"\"\n",
    "        数据集划分与网页一致：先分训练/验证集（val_rate），再从验证集分测试集（test_rate）\n",
    "        \"\"\"\n",
    "        # 划分训练集、验证集、测试集\n",
    "        X_train_val, X_test, y_train_val, y_test = train_test_split(\n",
    "            X, y, test_size=val_rate, random_state=102, stratify=y\n",
    "        )\n",
    "        X_val, X_test, y_val, y_test = train_test_split(\n",
    "            X_train_val, y_train_val, test_size=test_rate, random_state=102, stratify=y_train_val\n",
    "        )\n",
    "        \n",
    "        # 标准化（仅用训练集拟合）\n",
    "        scaler = StandardScaler()\n",
    "        X_train = scaler.fit_transform(X_train_val)  # 网页中使用全体训练验证集标准化？需确认\n",
    "        X_val = scaler.transform(X_val)\n",
    "        X_test = scaler.transform(X_test)\n",
    "        \n",
    "        # 根据mode选择数据\n",
    "        if mode == 'train':\n",
    "            self.data = (X_train, y_train_val)  # 网页中训练集是原始训练验证集（可能划分方式不同，需严格按网页代码）\n",
    "        elif mode == 'val':\n",
    "            self.data = (X_val, y_val)\n",
    "        elif mode == 'test':\n",
    "            self.data = (X_test, y_test)\n",
    "        else:\n",
    "            raise ValueError(\"mode must be 'train', 'val', or 'test'\")\n",
    "        \n",
    "        self.X, self.y = self.data\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.X)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return torch.tensor(self.X[idx], dtype=torch.float32), torch.tensor(self.y[idx], dtype=torch.long)\n",
    "\n",
    "# 3. 数据预处理管道（统一接口）\n",
    "def preprocess_data(data_path, time_steps=1024, window=128, batch_size=32):\n",
    "    \"\"\"\n",
    "    按网页逻辑预处理数据，返回DataLoader\n",
    "    \"\"\"\n",
    "    X, y = load_cwru_data(data_path, time_steps, window)\n",
    "    \n",
    "    # 创建数据集\n",
    "    train_dataset = CWRUDataset(X, y, mode='train', batch_size=batch_size)\n",
    "    val_dataset = CWRUDataset(X, y, mode='val', batch_size=batch_size)\n",
    "    test_dataset = CWRUDataset(X, y, mode='test', batch_size=batch_size)\n",
    "    \n",
    "    # 创建数据加载器\n",
    "    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "    \n",
    "    return train_loader, val_loader, test_loader\n",
    "\n",
    "# -------------------------- 模型定义（示例：保持PyTorch框架，可替换为用户原有模型） --------------------------\n",
    "class SimpleCNN(nn.Module):\n",
    "    def __init__(self, num_classes=10):\n",
    "        super(SimpleCNN, self).__init__()\n",
    "        self.conv_layers = nn.Sequential(\n",
    "            nn.Conv1d(1, 32, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm1d(32),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool1d(kernel_size=2),\n",
    "            \n",
    "            nn.Conv1d(32, 64, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool1d(kernel_size=2),\n",
    "            \n",
    "            nn.Conv1d(64, 128, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool1d(kernel_size=2)\n",
    "        )\n",
    "        self.fc_layers = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(128 * 130, 256),  # 需根据实际输出尺寸调整\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(256, num_classes)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = x.unsqueeze(1)  # 添加通道维度 [B, 1, L]\n",
    "        x = self.conv_layers(x)\n",
    "        x = self.fc_layers(x)\n",
    "        return x\n",
    "\n",
    "# -------------------------- 训练与评估（保持PyTorch风格） --------------------------\n",
    "def train_model(model, train_loader, val_loader, criterion, optimizer, device, epochs=50):\n",
    "    best_val_acc = 0.0\n",
    "    for epoch in range(epochs):\n",
    "        # 训练阶段\n",
    "        model.train()\n",
    "        total_loss, correct, total = 0.0, 0, 0\n",
    "        for inputs, targets in train_loader:\n",
    "            inputs, targets = inputs.to(device), targets.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, targets)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            _, pred = outputs.max(1)\n",
    "            correct += (pred == targets).sum().item()\n",
    "            total += targets.size(0)\n",
    "        train_acc = 100.0 * correct / total\n",
    "        \n",
    "        # 验证阶段\n",
    "        model.eval()\n",
    "        val_correct, val_total = 0, 0\n",
    "        with torch.no_grad():\n",
    "            for inputs, targets in val_loader:\n",
    "                inputs, targets = inputs.to(device), targets.to(device)\n",
    "                outputs = model(inputs)\n",
    "                _, pred = outputs.max(1)\n",
    "                val_correct += (pred == targets).sum().item()\n",
    "                val_total += targets.size(0)\n",
    "        val_acc = 100.0 * val_correct / val_total\n",
    "        \n",
    "        print(f\"Epoch {epoch+1}/{epochs} | Train Loss: {total_loss/total:.4f} | Train Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%\")\n",
    "        \n",
    "        # 保存最佳模型\n",
    "        if val_acc > best_val_acc:\n",
    "            best_val_acc = val_acc\n",
    "            torch.save(model.state_dict(), 'best_model.pth')\n",
    "    return model\n",
    "\n",
    "# -------------------------- 主函数（示例运行） --------------------------\n",
    "if __name__ == \"__main__\":\n",
    "    # 配置\n",
    "    data_path = './data/CWRU/'  # 替换为实际数据路径\n",
    "    time_steps = 1024\n",
    "    window = 128\n",
    "    batch_size = 32\n",
    "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "    \n",
    "    # 数据预处理\n",
    "    train_loader, val_loader, test_loader = preprocess_data(\n",
    "        data_path, time_steps=time_steps, window=window, batch_size=batch_size\n",
    "    )\n",
    "    \n",
    "    # 创建模型\n",
    "    model = SimpleCNN(num_classes=10).to(device)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "    \n",
    "    # 训练模型\n",
    "    model = train_model(model, train_loader, val_loader, criterion, optimizer, device, epochs=50)\n",
    "    \n",
    "    # 评估模型（可添加混淆矩阵等，此处略）\n",
    "    print(\"Training completed.\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
